from typing import Union

from ...executor.operator.physical_operator import CoveredIndexScan, HashAgg, IndexScan, NestedLoopJoin, PhysicalOperator, PhysicalQuery, Sort, TableScan
from ...common.fabric import FunctionColumn
from ...errors import SQLLogicalPlanError
from ..logical_operator import *
from ..parser.ast import *
from ...catalog import catalog_table, catalog_function, catalog_index

# 可用于判断表或者列是否存在，用于检验输入SQL语句的合法性
def table_exists(table_name: str):
    return len(catalog_table.select(lambda r: r.table_name == table_name)) > 0

def column_exists(table_name: str, column_name: str):
    return len(catalog_table.select(lambda r: r.table_name == table_name and column_name in r.columns)) > 0

def function_exists(function_name: str):
    results = catalog_function.select(lambda r: r.function_name == function_name and r.agg_function)
    return len(results) == 1

class SelectTransformer:
    @staticmethod
    def transform_clause_from(ast: Select, query: Query):
        # 表是不存在的，先验证一下，这个是逻辑计划生成的前提
        unchecked_tables = []
        if isinstance(ast.from_table, Identifier):
            unchecked_tables.append(ast.from_table.parts)
        elif isinstance(ast.from_table, Join):
            unchecked_tables.append(ast.from_table.left)
            unchecked_tables.append(ast.from_table.right)
        else:
            raise Exception("from_table is not a table or join")

        checked_tables = []
        for table_name in unchecked_tables:
            results = catalog_table.select(lambda r: r.table_name == table_name)
            
            if len(results) != 1:
                raise SQLLogicalPlanError(f'not found table {table_name}')
            checked_tables.append(table_name)
        
        for table_name in checked_tables:
            query.scan_operators.append(ScanOperator(table_name))

    @staticmethod
    def transform_target_list(ast: Select, query: Query):
        """处理投影部分"""
        for target in ast.targets:
            target_list = []
            if isinstance(target, Star):
                # 这个过程完成的是 * 的解析(将 * 转换为表的所有字段)
                for scan_operator in query.scan_operators:
                    table_name = scan_operator.table_name
                    results = catalog_table.select(lambda r: r.table_name == table_name)
                    # TableColumn 的小封装
                    for column in results[0].columns:
                        target_list.append(TableColumn(table_name, column))
            elif isinstance(target, Identifier):
                # 那就是表的一个具体列名
                # 一个小改进点：这个 target.parts 是可以包含 '.' 的，表示 t1.name 这种形式
                # 我们不考虑这种场景，要求用户必须指定表名，否则报错
                # 如果用户不指定表名，在 join 同名字段的时候，会报错，处理起来很麻烦
                full_name = target.parts
                if '.' not in full_name:
                    raise SQLLogicalPlanError(f'column {full_name} must be specified with table name')
                # t1.name -> t1, name
                # select t1.name from t1;
                # select name from t1; <-我们也可以自动来补充该t1表的信息，但是
                # 会增加代码量，其过程，就是从from后面所涉及的表中进行遍历列信息即可
                table_name, column = full_name.split('.')
                if not column_exists(table_name, column):
                    raise SQLLogicalPlanError(f'not found column {column} in table {table_name}')
                # pair 对，用一个class来封装，其实也是可以使用list或者tuple来封装
                # 但是，class封装会增加表达力，指的是我们可以通过判断数据类型，来知道是不是一个表的字段
                target_list.append(TableColumn(table_name, column))
            elif isinstance(target, FunctionOperation):
                # function 会比较麻烦一些，因为涉及到 function 内部还有些字段
                # e.g.，select count(t1.a) from t1 group by t1.a
                args = []
                for arg in target.args:
                    full_name = arg.parts
                    table_name, column = full_name.split('.')
                    if not column_exists(table_name, column):
                        raise SQLLogicalPlanError(f'not found column {column} in table {table_name}')
                    tc = TableColumn(table_name, column)
                    args.append(tc)
                function_name = target.op
                if not function_exists(function_name):
                    raise SQLLogicalPlanError(f'not found function {function_name}')
                target_list.append(FunctionColumn(function_name, *args))
            else:
                raise
            
            query.project_columns.extend(target_list)

    @staticmethod
    def transform_clause_where(ast: Select, query: Query):
        """处理选择部分"""
        if not ast.where:
            return
        # TODO: 暂时先不支持多个条件，先支持一个条件，后边有机会再增强
        query.where_condition = FilterOperator(Condition(ast.where))
        # 后面，把这个filter放到对应的scan算子的头上
    
    @staticmethod
    def transform_clause_join(ast: Select, query: Query):
        """处理连接部分"""
        # Join 语句出现在from_table字段中
        if not isinstance(ast.from_table, Join):
            return
        join_ast = ast.from_table

        # 做一个简单的检查，判断是否存在这个表
        # select t1.name, t2.age from t1 left join t2 on t1.id = t2.uid
        left_table_name = join_ast.left.parts
        right_table_name = join_ast.right.parts
        if not table_exists(left_table_name):
            raise SQLLogicalPlanError(f'not found table {left_table_name}')
        if not table_exists(right_table_name):
            raise SQLLogicalPlanError(f'not found table {right_table_name}')
        
        join_operator = JoinOperator(
                join_type=join_ast.join_type, 
                left_table_name=left_table_name, 
                right_table_name=right_table_name)

        query.join_operator = join_operator
        # 到这里还没完
        # 因为Join的子节点是数据的来源，而数据的来源是Scan算子，因此，我们为
        # JoinOperator 添加ScanOperator子节点
        for scan_operator in query.scan_operators:
            # 先匹配左节点，再匹配右节点，所以顺序不能变
            if scan_operator.table_name == left_table_name:
                join_operator.add_child(scan_operator)
            if scan_operator.table_name == right_table_name:
                join_operator.add_child(scan_operator)
            
        

    @staticmethod
    def transform_clause_order(ast: Select, query: Query):
        """处理排序部分"""
        if not ast.order_by:
            return
        full_name = ast.order_by.column
        if '.' not in full_name:
            raise SQLLogicalPlanError(f'column {full_name} must be specified with table name')
        table_name, column = full_name.split('.')

        # 检查下列是否存在
        if not column_exists(table_name, column):
            raise SQLLogicalPlanError(f'not found column {column} in table {table_name}')

        query.order_by = SortOperator(
            sort_column=TableColumn(table_name, column), 
            asc=(ast.order_by.direction == 'ASC')
        )


    @staticmethod
    def transform_clause_group(ast: Select, query: Query):
        """处理分组部分"""
        if not ast.group_by:
            return
        group_by_list = ast.group_by
        if len(group_by_list) != 1:
            raise NotImplementedError('not support multiple group by columns')

        # 检查下列是否存在
        full_name = ast.group_by[0].parts
        if '.' not in full_name:
            raise SQLLogicalPlanError(f'column {full_name} must be specified with table name')
        table_name, column = full_name.split('.')

        # 检查下列是否存在
        if not column_exists(table_name, column):
            raise SQLLogicalPlanError(f'not found column {column} in table {table_name}')

        query.group_by_column = TableColumn(table_name, column)

        # 检查是否存在聚合函数？在transform_target_list中已经完成了检查
        for column in query.project_columns:
            if isinstance(column, FunctionColumn) and catalog_function.select(lambda r: r.function_name == column.function_name and r.agg_function):
                query.aggregate_columns.append(column)
        
        if len(query.aggregate_columns) > 1:
            raise NotImplementedError('not support multiple aggregate functions')
        if query.aggregate_columns and len(query.aggregate_columns[0].args) != 1:
            raise NotImplementedError(f'aggregate function {query.aggregate_columns[0].function_name} not support multiple aggregate columns')
        

    @staticmethod
    def rewrite(query: Query):
        """逻辑计划的初步整理，优化"""
        # 我们组合各个query的属性，把他们按照顺序结构构建起一颗树
        # 这个顺序正好是 transform 过程反过来
        # all_seen_columns = {} # key: 表名，value: 列名
        building_node = query
        if query.group_by_column:
            operator = GroupOperator(
                group_by_column=query.group_by_column,
                aggregate_function_name=query.aggregate_columns[0].function_name,
                aggregate_column=query.aggregate_columns[0].args[0]
                )
            building_node = building_node.add_child(operator)
        if query.sort_operator:
            building_node = building_node.add_child(query.sort_operator)
        if query.join_operator:
            building_node = building_node.add_child(query.join_operator)
        if query.where_condition:
            # 把 where 条件下推到对应的扫描算子上
            # 在我们的场景中，有两种比较典型的场景不能直接推:
            # eg1: select * from t1, t2 where t1.a > t2.b;
            # eg2: select * from t1 where 1 > 2;
            # eg3: select * from t1 where t1.a > 100;
            # eg1 可以此时改写为 inner join
            # eg2 可以改写为直接返回空结果集
            # eg3 可以把 filter 条件下推到扫描过程中
            # 思考：你还能想到哪些不能推的场景呢？

            filter_operator = query.where_condition
            # eg1 的场景：
            if (isinstance(filter_operator.condition.left, TableColumn) and
                    isinstance(filter_operator.condition.right, TableColumn)):
                # 判断一下是否合法
                table_names = (query.join_operator.left_table_name,
                                query.join_operator.right_table_name)
                if not (query.join_operator and
                        filter_operator.condition.left.table_name in table_names and
                        filter_operator.condition.right.table_name in table_names):
                    raise SQLLogicalPlanError(f'tables in where clause should be all seen in '
                                                f'the join clause.')

                # 不支持的场景：
                # select * from t1 left join t2 on t1.id = t2.id where t1.age > t2.age;
                if query.join_operator.join_type != JoinType.CROSS_JOIN:
                    raise NotImplementedError('not supported complex where clause.')

                query.join_operator.join_type = JoinType.INNER_JOIN
                query.join_operator.condition = filter_operator.condition
            # 其他场景，这里面我们没有针对eg2单独改写，但是你可以试试~
            else:
                table_column = None
                if isinstance(filter_operator.condition.left, TableColumn):
                    table_column = filter_operator.condition.left
                elif isinstance(filter_operator.condition.right, TableColumn):
                    table_column = filter_operator.condition.right
                else:
                    # todo: 此处是 eg2 的改写地方
                    pass

                if table_column:
                    for scan_operator in query.scan_operators:
                        if scan_operator.table_name == table_column.table_name:
                            scan_operator.condition = filter_operator.condition

        if not query.join_operator:
            assert len(query.scan_operators) == 1
            building_node.add_child(query.scan_operators[0])
        # todo: 现在可以做一些RBO的优化动作了，例如列裁剪

    @staticmethod
    def transform(ast: Select) -> LogicalOperator:
        query = Query(Query.SELECT)
        SelectTransformer.transform_clause_from(ast, query)
        SelectTransformer.transform_target_list(ast, query)
        SelectTransformer.transform_clause_where(ast, query)
        SelectTransformer.transform_clause_join(ast, query)
        SelectTransformer.transform_clause_order(ast, query)
        SelectTransformer.transform_clause_group(ast, query)
        # 到此为止，我们就可以产生一个比较朴素的，原始的Query
        # 此时，我们可以进一步套用逻辑优化查询(重写)机制来优化Query
        SelectTransformer.rewrite(query)
        return query

def query_logical_plan(ast: ASTNode) -> LogicalOperator:
    if isinstance(ast, Select):
        return SelectTransformer.transform(ast)
    else:
        raise NotImplementedError('not supported query.')

# implementation 表示将逻辑计划转换为物理计划
# select 表示针对SELECT语句
class SelectImplementation:

    @staticmethod
    def implement_scan(node: ScanOperator) -> Union[TableScan, IndexScan, CoveredIndexScan]:
        # 返回的应该是 TableScan，IndexScan 或者 CoveredIndexScan 的一个
        # 我们此处定义一个rule：
        # 有索引，优先使用索引，如果这个索引可以用覆盖索引，那么优先用覆盖索引
        results = catalog_index.select(lambda r: r.table_name == node.table_name)
        if len(results) == 0 or not node.condition:
            # 没有索引，没有where条件 直接返回表扫描即可
            return TableScan(table_name=node.table_name, condition=node.condition)
        
        # 大于0，代表有索引，我们开始尝试使用索引
        # 我们来判断如下索引，能否使用
        # rule：假如有多个索引，我们选择列数最少的那个使用
        # 但是，这里面有个前提，是列能符合最左匹配原则
        candidate_indexes = []
        condition_columns = []
        if isinstance(node.condition.left, TableColumn):
            condition_columns.append(node.condition.left)
        if isinstance(node.condition.right, TableColumn):
            condition_columns.append(node.condition.right)

        # 这种边界场景：select t1.a from t1 where 1 > 2;
        if len(condition_columns) == 0:
            # 如果没有匹配上的，则直接返回表扫描
            return TableScan(table_name=node.table_name, condition=node.condition)
        
        # 下面，来找condition_columns与results中的索引列能够匹配上的
        # 如果匹配上了，则放到 candidate_indexes 中
        for catalog_index_form in results:
            # 下面这块逻辑就是在验证最左匹配原则(也就前缀匹配原则)，
            # 例如索引包含多个列 index(t1.a, t1.b, t1.c)
            # 而我们的谓词predicate在下面几个例子上，分别为：
            # 包含了 t1.a --> 符合
            # 包含了 t1.a, t1.b --> 符合
            # 包含了 t1.a, t1.b, t1.c --> 符合
            # 包含了 t1.b, t1.a --> 不符合
            matched = True
            for condition_column in condition_columns:
                for index_column in catalog_index_form.columns:
                    if condition_column.column_name != index_column:
                        matched = False
                        break
                if not matched:
                    break
            if matched:
                candidate_indexes.append(catalog_index_form)
        
        # 如果索引没有匹配上的，则直接返回表扫描
        if len(candidate_indexes) == 0:
            return TableScan(table_name=node.table_name, condition=node.condition)

        # 接下来，从这里面挑选最合适的一个索引
        # 如果，这个query中所涉及到的列，都涵盖在这个索引里面了，那么我们
        # 就使用覆盖索引，否则，我们使用普通索引，找在 candidate_indexes 索引列最短的那个使用
        # 因为索引越短，存储的冗余信息越少，IO效率就越高
        for candidate_index in candidate_indexes:
            if len(node.columns) == len(candidate_index.columns):
                # 此时，应该使用覆盖索引
                return CoveredIndexScan(
                    table_name=node.table_name,
                    index_name=candidate_index.index_name,
                    condition=node.condition,
                    columns=node.columns
                )

        # 如果没有覆盖索引，那么我们选择最短的那个索引
        shortest_index = min(candidate_indexes, key=lambda x: len(x.columns))
        return IndexScan(
            index_name=shortest_index.index_name,
            condition=node.condition
        )
        

    @staticmethod
    def implement_join(node: JoinOperator) -> NestedLoopJoin:
        # 可以优化的点：
        # 虽然说，我们支持两个节点进行Join，但是它们的顺序(内外表)，也是可以优化的点
        return NestedLoopJoin(
            join_type=node.join_type,
            left_table_name=node.left_table_name,
            right_table_name=node.right_table_name,
            join_condition=node.join_condition
        )

    @staticmethod
    def implement_sort(node: SortOperator) -> Sort:
        # 可以优化的点：
        # 默认是内排序，如果我们有参数设置的话，可以在此处做一个判断，如果超过了预估
        # 的阈值（内存的使用过大），那么就使用外排序
        return Sort(
            sort_column=node.sort_column,
            sort_order=node.asc
        )

    @staticmethod
    def implement_agg(node: GroupOperator) -> HashAgg:
        return HashAgg(
            group_by_column=node.group_by_column,
            aggregate_function_name=node.aggregate_function_name,
            aggregate_column=node.aggregate_column
        )
    
    @staticmethod
    def implement(node: ASTNode):
        # 这里是通过先递归逻辑计划树，然后产生
        # 与之对应的物理执行计划树
        if node is None:
            return None
        
        if isinstance(node, ScanOperator):
            physical_node = SelectImplementation.implement_scan(node)
        elif isinstance(node, JoinOperator):
            physical_node = SelectImplementation.implement_join(node)
        elif isinstance(node, SortOperator):
            physical_node = SelectImplementation.implement_sort(node)
        elif isinstance(node, GroupOperator):
            physical_node = SelectImplementation.implement_agg(node)
        elif isinstance(node, Query):
            physical_node = PhysicalQuery()
            # 此时，是最终要输出的所有列
            # 在这里，要对下面返回来的列，做裁剪（投影，也就是删除没有的列）
            physical_node.columns = node.project_columns
        else:
            raise NotImplementedError('not supported this type of node {node}.')
        
        # 没有子节点了
        if len(node.children) == 0:
            return physical_node

        # 还有子节点，相当于做了树的先序遍历
        for child in node.children:
            physical_node.add_child(SelectImplementation.implement(child))
        return physical_node


def query_physical_plan(logical_plan: LogicalOperator) -> PhysicalOperator:
    if isinstance(logical_plan, Query):
        if logical_plan.query_type == Query.SELECT:
            return SelectImplementation.implement(logical_plan)
        else:
            raise NotImplementedError(f'not supported this type of query {logical_plan.query_type}.')
    else:
        raise NotImplementedError(f'not supported this type of logical plan {logical_plan}.')

def query_plan(ast: ASTNode) -> PhysicalOperator:
    logical_plan = query_logical_plan(ast)
    return query_physical_plan(logical_plan)